#include "sockets.hpp"

#include "../luautil/luautil.hpp"



#include <gvl/sockets/sockets.hpp>
#include <gvl/sockets/fd_set.hpp>
#include <gvl/support/flags.hpp>
#include <gvl/support/bits.hpp>
#include <vector>
#include <deque>
#include <utility>
#include <memory>

#define LUASOCK(name) extern "C" int luasock_##name(lua_State* L)

LUASOCK(tcp_socket);
LUASOCK(udp_socket);
LUASOCK(select);

LUASOCK(address);
LUASOCK(address_from_name);

LUASOCK(connect);
LUASOCK(recv);
LUASOCK(send);
LUASOCK(sendto_array);
LUASOCK(readfrom_array);
LUASOCK(close);

static luaL_Reg sockets[] =
{
	{"tcp_socket", luasock_tcp_socket},
	{"udp_socket", luasock_udp_socket},
	{"new_address", luasock_address},
	{"select", luasock_select},

	{NULL, NULL}
};

static luaL_Reg sockets_socket[] =
{
	{"connect", luasock_connect},
	{"sendto_array", luasock_sendto_array},
	{"readfrom_array", luasock_readfrom_array},
	{"recv", luasock_recv},
	{"send", luasock_send},
	{"close", luasock_close},

	{NULL, NULL}
};

static luaL_Reg sockets_address[] =
{
	{"set_from_name", luasock_address_from_name},

	{NULL, NULL}
};

int sockets_socket_tag;
int sockets_address_tag;

int luaopen_sockets(lua_State* L)
{
	make_metatable(L, &sockets_socket_tag, sockets_socket, 0); // mt
	lua_pushvalue(L, -1); // ... | mt mt
	lua_setfield(L, -2, "__index"); // ... | mt
	lua_pop(L, 1); // <empty>

	make_metatable(L, &sockets_address_tag, sockets_address, 0); // mt
	lua_pushvalue(L, -1); // ... | mt mt
	lua_setfield(L, -2, "__index"); // ... | mt
	lua_pop(L, 1); // <empty>
	
	luaL_openlib(L, "sockets", sockets, 0); // sockets

	

	return 1;
}

struct luasock_socket
{
	luasock_socket(gvl::socket sock_init)
	: sock(sock_init)
	{
	}

	gvl::socket sock;

private:
	luasock_socket(luasock_socket const&);
	luasock_socket& operator=(luasock_socket const&);
};

inline luasock_socket& check_socket(lua_State* L, int idx)
{
	return check_object<luasock_socket>(L, idx, &sockets_socket_tag, "expected socket");
}

inline gvl::internet_addr& check_address(lua_State* L, int idx)
{
	return check_object<gvl::internet_addr>(L, idx, &sockets_address_tag, "expected inet address");
}

LUASOCK(tcp_socket)
{
	char const* addr = luaL_checkstring(L, 1);
	int port = luaL_checkint(L, 2);
	
	gvl::init_sockets();

	gvl::socket s = gvl::tcp_socket();
	if(!s.is_valid())
		return 0;
	s.set_nonblocking();
	
	void* mem = lua_newuserdata(L, sizeof(luasock_socket));
	luasock_socket* self = new (mem) luasock_socket(s);
	prepare_object(L, &sockets_socket_tag); // self -- self

	return 1;
}

LUASOCK(udp_socket)
{
	gvl::init_sockets();
	
	gvl::socket s = gvl::udp_socket();
	if(!s.is_valid())
		return 0;
	s.set_nonblocking();
	
	void* mem = lua_newuserdata(L, sizeof(luasock_socket));
	luasock_socket* self = new (mem) luasock_socket(s);
	prepare_object(L, &sockets_socket_tag); // self -- self

	return 1;
}

LUASOCK(address)
{
	void* mem = lua_newuserdata(L, sizeof(gvl::internet_addr));
	gvl::internet_addr* self = new (mem) gvl::internet_addr();
	prepare_object(L, &sockets_address_tag); // self -- self

	return 1;
}

LUASOCK(address_from_name)
{
	auto& self = check_address(L, 1);
	char const* addr = luaL_checkstring(L, 2);
	int port = luaL_checkint(L, 3);

	self = gvl::internet_addr(addr, port);
	return 0;
}

/* Convert a stack index to an absolute index. */
#define abs_index(L, i) \
  ((i) > 0 || (i) <= LUA_REGISTRYINDEX ? (i) : lua_gettop(L) + (i) + 1)

void table_to_fdset(lua_State* L, int tableidx, std::auto_ptr<gvl::select_fd_set>& fdset)
{
	if(lua_isnil(L, tableidx))
		return;
	fdset.reset(gvl::new_select_fd_set());
	luaL_argcheck(L, lua_istable(L, tableidx), tableidx, "expected table");

	for(int i = 1;; ++i)
	{
		lua_rawgeti(L, tableidx, i); // ... | socket?
		luasock_socket* socket = try_object<luasock_socket>(L, -1, &sockets_address_tag);
		if(socket)
			fdset->set(socket->sock);
		lua_pop(L, 1); // ... | <empty>
	}
}

// NOTE: This assumes orgtableidx is a table if fdset is non-zero!
void fdset_to_table(lua_State* L, int orgtableidx, std::auto_ptr<gvl::select_fd_set>& fdset)
{
	orgtableidx = abs_index(L, orgtableidx);

	lua_newtable(L); // rettable
	if(!fdset.get())
		return; // Empty table

	int desti = 1;
	for(int i = 1;; ++i)
	{
		lua_rawgeti(L, orgtableidx, i); // rettable socket?
		luasock_socket* socket = try_object<luasock_socket>(L, -1, &sockets_address_tag);
		if(socket && fdset->is_set(socket->sock))
		{
			lua_pushvalue(L, -1); // rettable socket socket
			lua_rawseti(L, -3, desti); // rettable(-3) socket(-2) socket(-1) -- rettable socket
			++desti;
			lua_pushvalue(L, -1); // rettable socket socket
			lua_rawset(L, -3); // rettable(-3) socket(-2) socket(-1) | rettable
		}
		else
		{
			lua_pop(L, 1); // rettable socket? -- rettable
		}
	}
}

LUASOCK(select)
{
	auto& self = check_socket(L, 1);
	std::auto_ptr<gvl::select_fd_set> fdset[3];
	//int timeout = luaL_checkint(L, 2);
	table_to_fdset(L, 2, fdset[0]);
	table_to_fdset(L, 3, fdset[1]);
	table_to_fdset(L, 4, fdset[2]);

	int r = gvl::select(fdset[0].get(), fdset[1].get(), fdset[2].get());

	fdset_to_table(L, 2, fdset[0]); // readtab
	fdset_to_table(L, 3, fdset[1]); // readtab writetab
	fdset_to_table(L, 4, fdset[2]); // readtab writetab errtab
	lua_pushinteger(L, r); // readtab writetab errtab ret
	return 4;
}

LUASOCK(connect)
{
	auto& self = check_socket(L, 1);
	auto& addr = check_address(L, 2);

	int r = self.sock.connect(addr);
	lua_pushinteger(L, r);
	return 1;
}

LUASOCK(bind)
{
	auto& self = check_socket(L, 1);
	int port = luaL_optint(L, 2, 0);

	int r = self.sock.bind(port);
	lua_pushinteger(L, r);
	return 1;
}

LUASOCK(recv)
{
	auto& self = check_socket(L, 1);

	char buffer[4096];
	int r = self.sock.recv(buffer, 4096);

	lua_pushinteger(L, r);

	if (r == 0)
	{
		return 1;
	}
	else if (r == gvl::socket::would_block)
	{
		lua_pushlstring(L, "", 0); // Empty string
		return 2;
	}
	else if (r < 0)
	{
		return 1;
	}

	lua_pushlstring(L, buffer, r);
	return 2;
}

LUASOCK(send)
{
	auto& self = check_socket(L, 1);
	std::size_t len;
	char const* data = luaL_checklstring(L, 2, &len);
	int offset = luaL_optint(L, 3, 0);

	if(offset > len)
	{
		lua_pushinteger(L, 0);
	}
	else
	{
		int r = self.sock.send(data + offset, len - offset);
		lua_pushinteger(L, r);
	}
	return 1;
}

int const max_udp_packet_size = 2048;

LUASOCK(sendto_array)
{
	auto& self = check_socket(L, 1);
	auto& addr = check_address(L, 2);
	luaL_argcheck(L, lua_istable(L, 3), 3, "expected table");
	int len = luaL_checkint(L, 4);

	if(len < 0 || len > max_udp_packet_size)
		return 0;

	int words = (len + 3) >> 2;

	uint32_t buffer[(max_udp_packet_size + 3) / 4];
	double* arr = 1 + static_cast<double*>(luaext_rawarraybase(L, 3, words));
	
	for(int i = 0; i < words; ++i)
		buffer[i] = gvl::bswap(gvl::int32_as_uint32(int32_t(arr[i])));

	int r = self.sock.sendto(buffer, len, addr);
	lua_pushinteger(L, r);
	return 1;
}

LUASOCK(readfrom_array)
{
	auto& self = check_socket(L, 1);
	auto& addr = check_address(L, 2);
	luaL_argcheck(L, lua_istable(L, 3), 3, "expected table");

	uint32_t buffer[(max_udp_packet_size + 3) / 4] = {};

	int r = self.sock.recvfrom(buffer, max_udp_packet_size, addr);

	if(r >= 0)
	{
		int words = (r + 3) >> 2;
		double* arr = 1 + static_cast<double*>(luaext_rawarraybase(L, 3, words));

		for(int i = 0; i < words; ++i)
			arr[i] = gvl::uint32_as_int32(gvl::bswap(buffer[i]));
	}

	lua_pushinteger(L, r);
	return 1;
}

LUASOCK(close)
{
	auto& self = check_socket(L, 1);

	self.sock.close();
	self.sock = gvl::socket::invalid();
	return 0;
}
